from survtrace.dataset_own import load_data
from survtrace.evaluate_utils import Evaluator
from survtrace.utils import set_random_seed
from survtrace.model import SurvTraceSingle
from survtrace.train_utils import Trainer
from survtrace.config import STConfig

import matplotlib.pyplot as plt
# define the setup parameters
STConfig['data'] = 'seer'
STConfig['num_hidden_layers'] = 4
STConfig['hidden_size'] = 16
STConfig['intermediate_size'] = 64
STConfig['num_attention_heads'] = 8
STConfig['initializer_range'] = .02
STConfig['early_stop_patience'] = 5
set_random_seed(STConfig['seed'])

hparams = {
    'batch_size': 512,
    'weight_decay': 0,
    'learning_rate': 1e-5,
    # 'learning_rate': 0.0001,
    'epochs': 1000,
}
# load data
df, df_train, df_y_train, df_test, df_y_test, df_val, df_y_val = load_data(STConfig)

# get model
model = SurvTraceSingle(STConfig)
print(model)

# initialize a trainer & start training
trainer = Trainer(model)
train_loss_list, val_loss_list = trainer.fit((df_train, df_y_train), (df_val, df_y_val), (df_test, df_y_test),
        batch_size=hparams['batch_size'],
        epochs=hparams['epochs'],
        learning_rate=hparams['learning_rate'],
        weight_decay=hparams['weight_decay'],
)

# evaluate model
# evaluator = Evaluator(df, df_train.index)
evaluator = Evaluator()
evaluator.eval(model, (df_test, df_y_test))
print("done")

plt.plot(train_loss_list, label='train')
plt.plot(val_loss_list, label='val')
plt.legend(fontsize=20)
plt.xlabel('epoch',fontsize=20)
plt.ylabel('loss', fontsize=20)
plt.show()